import torch
from tqdm import tqdm
from torch_geometric.data import Data, InMemoryDataset
from sklearn.preprocessing import LabelEncoder
import pandas as pd

df = pd.read_csv("data/wikitext2.csv")
df.columns = ["sid", "word_id", "category", "label"]
# print(df)

# print(df[:5])
# from torch_geometric.data import InMemoryDataset
# 测试预处理训练集合信息
df_test = df[:2000]

# 使用ｓｉｄ进行ｇｒｏｕｐｂｙ操作并创建一个迭代器，每次迭代出一个ｓｉｄ和一组数据
grouped = df_test.groupby("sid")
# print(grouped)
for sid, group in tqdm(grouped):
    word_local_index = LabelEncoder().fit_transform(group.word_id)
    print("word_local_indexword_local_indexword_local_index",word_local_index)
    group = group.reset_index(drop=True)
    group["word_local_index"] = word_local_index
    print(sid, group)
    node_features = group.loc[group.sid == sid, ["word_local_index", "word_id"]].sort_values(
        "word_local_index").word_id.drop_duplicates().values
    print("node_features", node_features)

    node_features = torch.LongTensor(node_features).unsqueeze(1)
    target_nodes = group.word_local_index.values[1:]
    source_nodes = group.word_local_index.values[:-1]
    print("source_nodes",source_nodes)
    print("target_nodes",target_nodes)
    edge_index = torch.tensor([source_nodes, target_nodes], dtype=torch.long)
    x = node_features
    y = torch.FloatTensor([group.label.values[0]])
    data = Data(x=x, edge_index=edge_index, y=y)
    print("data", data)
